/*
 * This program was designed for pair-end reads.
 * After bwa mapping, there is a need to filter those not properly paired templates
 * like singleton, or discordant mapping, or multiple hits
 *
 * About the defination of unique mapping, there are different view:
 * 1. just as bwa defination: only one best hit (X0 was 1, X1 was 0)
 * 2. but secondary hits sometimes are small drop than best hit, the another thought is that taking X1 into consideration
 *
 * There are several strategies when filtering
 * 1. just keep those unique mapping fragments
 * 2. keep those unique mapping fragments and proper paired templates
 * 3. unique & proper mapping at template level
 *
 * So, this script provides 2*3 strategies for filtering
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "sam.h"

#define N_REPORT 500000
#define PROG "bam.unique"
#define bam1_rid(b) (((b)->core.flag & BAM_FREAD1)? 1: 2)
#define bam1_unmapped(b) ((b)->core.flag & BAM_FUNMAP)
#define bam1_proper(b) ((b)->core.flag & BAM_FPROPER_PAIR)
#define flag_against_coor(b) (((b)->core.tid < 0) ^ ((b)->core.flag & BAM_FUNMAP))
#define is_proper_tpl(p) ((p)[0]->core.flag & (p)[1]->core.flag & BAM_FPROPER_PAIR)
#define span_diff_chrom(p) ((p)[0]->core.tid != (p)[1]->core.tid)
#define has_xa(b) bam_aux_get((b), "XA")

typedef int (*strategy)(bam1_t **, int *);
typedef int boolean;

int usage(void);
int readpair(samfile_t *, bam1_t **);
void tpl_sort(bam1_t **);
boolean is_unique(bam1_t *);
boolean is_unique2(bam1_t *);
int contain_sw(bam1_t **, int *);
int contain_unique(bam1_t **, int *);
int contain_unique2(bam1_t **, int *);
int contain_unmapped(bam1_t **, int *);

int keep_unique(bam1_t **, int *);
int keep_unique2(bam1_t **, int *);
int keep_unique_and_proper(bam1_t **, int *);
int keep_unique2_and_proper(bam1_t **, int *);
int keep_unique_tpl(bam1_t **, int *);
int keep_unique2_tpl(bam1_t **, int *);

#define N_MAT 26
enum {
	tpl_multiple_frag, tpl_single_frag,
	tpl_proper,
		tpl_bridge_diff_chrom,
		tpl_sw,
			tpl_sw_discordant,
				tpl_sw_discordant_mU, tpl_sw_discordant_mU2,
			tpl_sw_singleton,
				tpl_sw_singleton_mU, tpl_sw_singleton_mU2,
		tpl_aln,
			tpl_aln_UU, tpl_aln_DU, tpl_aln_DD,
			tpl_aln_UU2, tpl_aln_DU2, tpl_aln_DD2,
	tpl_both_no_map,
	tpl_singleton,
		tpl_singleton_U, tpl_singleton_D,
		tpl_singleton_U2, tpl_singleton_D2,
	tpl_span_chrom,
	tpl_other_discordant
} MAT;

int main (int argc, char *argv[]){
	char *fn_matrics = 0, *fn_log = 0, *method = 0;
	int is_bamin = 1, is_bamout = 1, with_header = 0, compress_level = -1, expand_unique_def = 0;
	int c = -1, which;
	unsigned long mat[N_MAT] = {0ul};
	strategy pass = 0;

	// handle commond-line arguments
	while((c = getopt(argc, argv, "?Suhset:m:l:?")) >= 0){
		switch(c){
			case 'S': is_bamin = 0; break;
			case 'u': compress_level = 0; break;
			case 'h': with_header = 1;
			case 's': is_bamout = 0; break;
			case 'e': expand_unique_def = 1; break;
			case 't': method = strdup(optarg); break;
			case 'm': fn_matrics = strdup(optarg); break;
			case 'l': fn_log = strdup(optarg); break;
			case '?':
			default: return usage();
		}
	}
	if(argc != optind + 2)	return usage(); // potential memory leak ...

	// set BAM/SAM read/write mode
	char in_mode[5], out_mode[5];
	strcpy(in_mode, "r"); strcpy(out_mode, "w");
	if(is_bamin)	strcat(in_mode, "b");
	if(is_bamout){
		strcat(out_mode, "b");
		if(compress_level >= 0){
			char tmp[2];
			tmp[0] = compress_level + '0'; tmp[1] = '\0';
			strcat(out_mode, tmp);
		}
	}else if(with_header)	strcat(out_mode, "h");
	// open BAM/SAM stream
	samfile_t *in = 0, *out = 0;
	if((in = samopen(argv[optind], in_mode, 0)) == 0){
		fprintf(stderr, "[%s] fatal error: fail to open \"%s\" for reading.\n", PROG, argv[optind]);
		return 1;
	}
	optind ++;
	if((out = samopen(argv[optind], out_mode, in->header)) == 0){
		fprintf(stderr, "[%s] fatal error: fail to create \"%s\" for writing.\n", PROG, argv[optind]);
		return 1;
	}
	FILE *LOG = stderr;
	// open log file
	if(fn_log)
		if((LOG = fopen(fn_log, "w")) == 0){
			fprintf(stderr, "[%s] fail to create log file \"%s\".\n", PROG, fn_log);
			LOG = stderr;
		}

	// assign strategy
	if(method == 0)	pass = keep_unique_tpl;
	else if(strcmp("keep_unique", method) == 0){
		if(expand_unique_def)	pass = keep_unique2;
		else	pass = keep_unique;
	}else if(strcmp("keep_unique_and_proper", method) == 0){
		if(expand_unique_def)	pass = keep_unique2_and_proper;
		else	pass = keep_unique_and_proper;
	}else if(strcmp("keep_unique_tpl", method) == 0){
		if(expand_unique_def)	pass = keep_unique2_tpl;
		else	pass = keep_unique_tpl;
	}else	pass = keep_unique_tpl;
	// init scan
	bam1_t * pair[2];
	c = 0;
	long n = 0, j = 1;
	while((c = readpair(in, pair)) > 0){ // read a paired-end
		n += c;
		if(n >= j * N_REPORT){
			j ++;
			fprintf(stderr, "[%s] processed %ld reads.\n", PROG, n);
		}
		if(c > 2){
			mat[tpl_multiple_frag] ++;
			fprintf(LOG, "%s pMR\n", bam1_qname(pair[0]));
		}else if(c == 1){
			mat[tpl_single_frag] ++;
			fprintf(LOG, "%s pSR\n", bam1_qname(pair[0]));
		}else{ // only care about pair end
			tpl_sort(pair);
			if(is_proper_tpl(pair)){ // proper paired
				mat[tpl_proper] ++;
				if(flag_against_coor(pair[0]) || flag_against_coor(pair[1])){
					// because of bridge different chromosome
					mat[tpl_bridge_diff_chrom] ++;
				}else if(contain_sw(pair, &which)){ // mated or fixed by SW
					mat[tpl_sw] ++;
					if(has_xa(pair[which - 1])){ // discordant
						mat[tpl_sw_discordant] ++;
						if(is_unique(pair[2-which]))
							mat[tpl_sw_discordant_mU] ++;
						if(is_unique2(pair[2-which]))
							mat[tpl_sw_discordant_mU2] ++;
					}else{ // singleton
						mat[tpl_sw_singleton] ++;
						if(is_unique(pair[2-which]))
							mat[tpl_sw_singleton_mU] ++;
						if(is_unique2(pair[2-which]))
							mat[tpl_sw_singleton_mU2] ++;
					}
				}else{ // both end mapped previously
					mat[tpl_aln] ++;
					if(contain_unique(pair, &which)){
						if(which == 3) // both unique
							mat[tpl_aln_UU] ++;
						else	mat[tpl_aln_DU] ++;
					}else // both multiple
						mat[tpl_aln_DD] ++;
					if(contain_unique2(pair, &which)){
						if(which == 3) // both unique
							mat[tpl_aln_UU2] ++;
						else	mat[tpl_aln_DU2] ++;
					}else // both multiple
						mat[tpl_aln_DD2] ++;
				}
			}else if(contain_unmapped(pair, &which)){ // exist unmapped end
				if(which == 3)	mat[tpl_both_no_map] ++;
				else{ // only one
					mat[tpl_singleton] ++;
					if(is_unique(pair[2-which]))
						mat[tpl_singleton_U] ++;
					else	mat[tpl_singleton_D] ++;
					if(is_unique2(pair[2-which]))
						mat[tpl_singleton_U2] ++;
					else	mat[tpl_singleton_D2] ++;
				}
			}else{ // other discordant reason
				if(span_diff_chrom(pair))
					mat[tpl_span_chrom] ++;
				else	mat[tpl_other_discordant] ++;
			}
			// print out
			pass(pair, &which);
			if(which & 1)	samwrite(out, pair[0]);
			if(which & 2)	samwrite(out, pair[1]);
		}
	}

	samclose(in);
	samclose(out);

	FILE *MATRICS = stderr;
	if(fn_matrics){
		if((MATRICS = fopen(fn_matrics, "w")) == 0){
			fprintf(stderr, "[%s] fail to create file \"%s\" for writing.\n", PROG, fn_matrics);
			return 1;
		}
	}
	{
		fprintf(MATRICS, "##matrics\n");
		fprintf(MATRICS, "%-25s\t%-25s\n%-25ld\t%-25ld\n", "tpl_multiple_frag", "tpl_single_frag", mat[tpl_multiple_frag], mat[tpl_single_frag]);
		fprintf(MATRICS, "%-25s\n%-25ld\n", "tpl_proper", mat[tpl_proper]);
		fprintf(MATRICS, "\t%-25s\n\t%-25ld\n", "tpl_bridge_diff_chrom", mat[tpl_bridge_diff_chrom]);
		fprintf(MATRICS, "\t%-25s\n\t%-25ld\n", "tpl_sw", mat[tpl_sw]);
		fprintf(MATRICS, "\t\t%-25s\n\t\t%-25ld\n", "tpl_sw_discordant", mat[tpl_sw_discordant]);
		fprintf(MATRICS, "\t\t\t%-25s\t%-25s\n\t\t\t%-25ld\t%-25ld\n", "tpl_sw_discordant_mU", "tpl_sw_discordant_mU2", mat[tpl_sw_discordant_mU], mat[tpl_sw_discordant_mU2]);
		fprintf(MATRICS, "\t\t%-25s\n\t\t%-25ld\n", "tpl_sw_singleton", mat[tpl_sw_singleton]);
		fprintf(MATRICS, "\t\t\t%-25s\t%-25s\n\t\t\t%-25ld\t%-25ld\n", "tpl_sw_singleton_mU", "tpl_sw_singleton_mU2", mat[tpl_sw_singleton_mU], mat[tpl_sw_singleton_mU2]);
		fprintf(MATRICS, "\t%-25s\n\t%-25ld\n", "tpl_aln", mat[tpl_aln]);
		fprintf(MATRICS, "\t\t%-25s\t%-25s\t%-25s\n\t\t%-25ld\t%-25ld\t%-25ld\n", "tpl_aln_UU", "tpl_aln_DU", "tpl_aln_DD", mat[tpl_aln_UU], mat[tpl_aln_DU], mat[tpl_aln_DD]);
		fprintf(MATRICS, "\t\t%-25s\t%-25s\t%-25s\n\t\t%-25ld\t%-25ld\t%-25ld\n", "tpl_aln_UU2", "tpl_aln_DU2", "tpl_aln_DD2", mat[tpl_aln_UU2], mat[tpl_aln_DU2], mat[tpl_aln_DD2]);
		fprintf(MATRICS, "%-25s\n%-25ld\n", "tpl_both_no_map", mat[tpl_both_no_map]);
		fprintf(MATRICS, "%-25s\n%-25ld\n", "tpl_singleton", mat[tpl_singleton]);
		fprintf(MATRICS, "\t%-25s\t%-25s\n\t%-25ld\t%-25ld\n", "tpl_singleton_U", "tpl_singleton_D", mat[tpl_singleton_U], mat[tpl_singleton_D]);
		fprintf(MATRICS, "\t%-25s\t%-25s\n\t%-25ld\t%-25ld\n", "tpl_singleton_U2", "tpl_singleton_D2", mat[tpl_singleton_U2], mat[tpl_singleton_D2]);
		fprintf(MATRICS, "%-25s\n%-25ld\n", "tpl_span_chrom", mat[tpl_span_chrom]);
		fprintf(MATRICS, "%-25s\n%-25ld\n", "tpl_other_discordant", mat[tpl_other_discordant]);
		fclose(MATRICS);
	}
	fclose(LOG);
	return EXIT_SUCCESS;
}

/*
 * @dscp: read reads pair from name-sorted SAM/BAM file. Assume one template appears at most twice.
 * @para f: SAM/BAM stream
 * @para pair: reference carrier for tempalte alignment
 * @return: 0 for EOF, 1 for single mapping, 2 for pair mapping, >=3 for multiple mapping
 */
int readpair (samfile_t *f, bam1_t *pair[2]){
	static bam1_t *loop[3] = {0, 0, 0};
	static int b = 0, e = 0;
	int m = 2, c = -1;

	pair[0] = (bam1_t *)(0); pair[1] = (bam1_t *)(0);

	if(loop[b] == 0){ // only run once, init buckets                       
		for(e = 0; e < 3; e ++)	loop[e] = bam_init1();
		e = b;
	}
	m = 2;
	while((c = samread(f, loop[e])) >= 0){
		if(strcmp(bam1_qname(loop[e]), bam1_qname(loop[b])) == 0){
			if((e-b+3)%3 == 2)	m ++;// keep e unchanging
			else	e = (1+e)%3;
		}else	break;
	}
	if(c < 0){ // meet the EOF
		if(b == e){
			for(e = 0; e < 3; e ++)	bam_destroy1(loop[e]);
			return 0; // cache has been readout
		}else{
			pair[0] = loop[b];
			b = (1+b)%3;
			if(b != e){ pair[1] = loop[b]; b = e; }
		}
	}else{ // meet another template
		pair[0] = loop[b];
		b = (1+b)%3;
		if(b != e){ pair[1] = loop[b]; b = e; }
		e = (1+e)%3;
	}
	return pair[1] == 0? 1: m;
}

int x0 (bam1_t *b){
	int n;
	if(bam_aux_get(b, "X0")){
		n = bam_aux2i(bam_aux_get(b, "X0"));
		if(n > 1)	return 2;
		else if(n == 1)	return 1;
		else	return 0;
	}else return 0;
}
boolean x1 (bam1_t *b){
	int n;
	if(bam_aux_get(b, "X1")){
		n = bam_aux2i(bam_aux_get(b, "X1"));
		return n>0? 1: 0;
	}else return 0;
}
char xt (bam1_t *b){
	char c = 0;
	if(bam_aux_get(b, "XT"))
		c = bam_aux2A(bam_aux_get(b, "XT"));
	return c;
}
void tpl_sort (bam1_t *p[2]){
	if(bam1_rid(p[0]) != 1){
		bam1_t *tmp = p[0];
		p[0] = p[1];
		p[1] = tmp;
	}
}
boolean is_unique (bam1_t *b){
	return (xt(b) == 'U');
}
boolean is_unique2 (bam1_t *b){
	return ((xt(b) == 'U') && x1(b) == 0);
}
int contain_sw(bam1_t *p[2], int *w){
	*w = 0;
	if(xt(p[0]) == 'M')	(*w) |= 1;
	if(xt(p[1]) == 'M')	(*w) |= 2;
	return *w;
}
int contain_unique(bam1_t *p[2], int *w){
	*w = 0;
	if(is_unique(p[0]))	(*w) |= 1;
	if(is_unique(p[1]))	(*w) |= 2;
	return *w;
}
int contain_unique2(bam1_t *p[2], int *w){
	*w = 0;
	if(is_unique2(p[0]))	(*w) |= 1;
	if(is_unique2(p[1]))	(*w) |= 2;
	return *w;
}
int contain_unmapped(bam1_t *p[2], int *w){
	*w = 0;
	if(bam1_unmapped(p[0]))	(*w) |= 1;
	if(bam1_unmapped(p[1]))	(*w) |= 2;
	return *w;
}

int keep_unique(bam1_t *p[2], int *w){
	return contain_unique(p, w);
}
int keep_unique2(bam1_t *p[2], int *w){
	return contain_unique2(p, w);
}
int keep_unique_and_proper(bam1_t *p[2], int *w){
	*w = 0;
	if(is_unique(p[0]) && bam1_proper(p[0]))	(*w) |= 1;
	if(is_unique(p[1]) && bam1_proper(p[1]))	(*w) |= 2;
	return *w;
}
int keep_unique2_and_proper(bam1_t *p[2], int *w){
	*w = 0;
	if(is_unique2(p[0]) && bam1_proper(p[0]))	(*w) |= 1;
	if(is_unique2(p[1]) && bam1_proper(p[1]))	(*w) |= 2;
	return *w;
}
int keep_unique_tpl(bam1_t *p[2], int *w){
	*w = 0;
	if(!(bam1_unmapped(p[0]) || bam1_unmapped(p[1])) && bam1_proper(p[0]) && bam1_proper(p[1]) &&
		(is_unique(p[0]) || is_unique(p[1])))
		*w = 3;
	return *w;
}
int keep_unique2_tpl(bam1_t *p[2], int *w){
	*w = 0;
	if(!(bam1_unmapped(p[0]) || bam1_unmapped(p[1])) && bam1_proper(p[0]) && bam1_proper(p[1]) &&
		(is_unique2(p[0]) || is_unique2(p[1])))
		*w = 3;
	return *w;
}

int usage (void){
	fprintf(stderr, "\n");
	fprintf(stderr, "Usage:    %s [options] <in.bam> <out.bam>\n\n", PROG);
	fprintf(stderr, "Options:  -S         Input is SAM. Default is BAM\n");
	fprintf(stderr, "          -s         Output is SAM. Default is BAM\n");
	fprintf(stderr, "          -u         Output is uncompressed BAM for pipe. Default is BAM, where all unique reads including paired-end or singleton reads are written to\n");
	fprintf(stderr, "          -h         Output is SAM with header\n");
	fprintf(stderr, "          -e         expand unique defination which taking X1 into consideration\n");
	fprintf(stderr, "          -t method  choose filter method: keep_unique; keep_unique_and_proper; keep_unique_tpl (default)\n");
	fprintf(stderr, "          -m file    Matrics file records count\n");
	fprintf(stderr, "          -l log     Log to log file\n");
	fprintf(stderr, "          -?         Show this help message\n\n");
	return 1;
}
